import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

from ultralytics import YOLO
from pathlib import Path

def train_model():
    # 使用本地的yolo11n-pose.pt权重进行训练
    weights_path = Path('yolo11n-pose.pt')
    if not weights_path.exists():
        raise FileNotFoundError(f"权重文件 {weights_path} 不存在！")
    
    # 使用官方推荐的方式加载预训练模型
    model = YOLO("yolo11n-pose.yaml").load(str(weights_path.absolute()))
    
    # 训练配置
    results = model.train(
        data='dataset/dataset.yaml',    # 数据集配置文件
        epochs=100,                     # 训练轮数
        imgsz=640,                      # 输入图像尺寸
        batch=16,                       # 批次大小
        workers=8,                      # 数据加载线程数
        device=0,                       # 使用GPU训练
        project='runs/train',           # 项目目录
        name='yolov8n_pose',           # 实验名称
        exist_ok=True,                  # 允许覆盖已存在的实验目录
        pretrained=False,               # 不使用在线预训练权重
        optimizer='auto',               # 优化器选择
        verbose=True,                   # 显示详细信息
        seed=42,                        # 随机种子
        patience=50,                    # 早停耐心值
        save=True,                      # 保存模型
        save_period=-1,                 # 每个epoch都保存
        val=True,                       # 进行验证
        amp=True,                       # 使用混合精度训练
        resume=False                    # 不恢复之前的训练
    )
    
    # 验证最终模型
    metrics = model.val()
    print(f"验证指标:\n{metrics}")
    
    # 导出模型
    path = model.export(format='onnx')  # 导出为ONNX格式
    print(f"模型已导出到: {path}")
    
    return results

if __name__ == "__main__":
    # 创建必要的目录
    Path('runs/train').mkdir(parents=True, exist_ok=True)
    
    # 开始训练
    results = train_model()
    
    # 打印训练结果
    print("\n训练完成！")
    print(f"最佳mAP: {results.metrics.get('metrics/mAP50-95(P)', 0)}")
    print(f"最佳epoch: {results.best_epoch}") 